"""
交易日历模块
判断是否是交易日，获取最近交易日
"""

import datetime
import pandas as pd
from typing import Optional


class TradingCalendar:
    """交易日历类"""
    
    # 固定节假日（不包括周末）
    FIXED_HOLIDAYS = {
        (1, 1): "元旦",
        (5, 1): "劳动节",
        (10, 1): "国庆节",
        (10, 2): "国庆节",
        (10, 3): "国庆节",
    }
    
    # 2025年节假日（包括调休）
    HOLIDAYS_2025 = [
        # 元旦
        datetime.date(2025, 1, 1),
        # 春节
        datetime.date(2025, 1, 28),
        datetime.date(2025, 1, 29),
        datetime.date(2025, 1, 30),
        datetime.date(2025, 1, 31),
        datetime.date(2025, 2, 1),
        datetime.date(2025, 2, 2),
        datetime.date(2025, 2, 3),
        datetime.date(2025, 2, 4),
        # 清明节
        datetime.date(2025, 4, 4),
        datetime.date(2025, 4, 5),
        datetime.date(2025, 4, 6),
        # 劳动节
        datetime.date(2025, 5, 1),
        datetime.date(2025, 5, 2),
        datetime.date(2025, 5, 3),
        datetime.date(2025, 5, 4),
        datetime.date(2025, 5, 5),
        # 端午节
        datetime.date(2025, 5, 31),
        datetime.date(2025, 6, 1),
        datetime.date(2025, 6, 2),
        # 中秋节
        datetime.date(2025, 10, 6),
        datetime.date(2025, 10, 7),
        datetime.date(2025, 10, 8),
        # 国庆节
        datetime.date(2025, 10, 1),
        datetime.date(2025, 10, 2),
        datetime.date(2025, 10, 3),
        datetime.date(2025, 10, 4),
        datetime.date(2025, 10, 5),
    ]
    
    @staticmethod
    def is_trading_day(date: Optional[datetime.date] = None) -> bool:
        """
        判断是否是交易日
        
        Args:
            date: 要判断的日期，默认为今天
            
        Returns:
            bool: True为交易日，False为非交易日
        """
        if date is None:
            date = datetime.date.today()
        
        # 检查是否是周末
        if date.weekday() >= 5:  # 5=周六, 6=周日
            return False
        
        # 检查是否在节假日列表中
        if date in TradingCalendar.HOLIDAYS_2025:
            return False
        
        return True
    
    @staticmethod
    def get_last_trading_day(date: Optional[datetime.date] = None) -> datetime.date:
        """
        获取最近的交易日
        
        Args:
            date: 基准日期，默认为今天
            
        Returns:
            datetime.date: 最近的交易日
        """
        if date is None:
            date = datetime.date.today()
        
        # 如果当天就是交易日，直接返回
        if TradingCalendar.is_trading_day(date):
            return date
        
        # 往前查找最近的交易日（最多查找15天）
        for i in range(1, 16):
            check_date = date - datetime.timedelta(days=i)
            if TradingCalendar.is_trading_day(check_date):
                return check_date
        
        # 如果15天内都没找到，返回当前日期（异常情况）
        return date
    
    @staticmethod
    def get_trading_status() -> dict:
        """
        获取当前交易状态
        
        Returns:
            dict: 包含交易状态信息的字典
        """
        today = datetime.date.today()
        now = datetime.datetime.now()
        
        is_trading = TradingCalendar.is_trading_day(today)
        last_trading = TradingCalendar.get_last_trading_day(today)
        
        # 判断是否在交易时间内（9:30-15:00）
        market_open_time = datetime.time(9, 30)
        market_close_time = datetime.time(15, 0)
        current_time = now.time()
        
        is_trading_hours = False
        if is_trading:
            # 检查是否在午休时间（11:30-13:00）
            lunch_start = datetime.time(11, 30)
            lunch_end = datetime.time(13, 0)
            
            if lunch_start <= current_time < lunch_end:
                is_trading_hours = False  # 午休
                market_status = "午休"
            elif market_open_time <= current_time < market_close_time:
                is_trading_hours = True
                market_status = "交易中"
            elif current_time < market_open_time:
                market_status = "未开盘"
            else:
                market_status = "已收盘"
        else:
            market_status = "休市"
        
        return {
            '今日日期': today.strftime('%Y-%m-%d'),
            '是否交易日': is_trading,
            '市场状态': market_status,
            '是否交易时间': is_trading_hours,
            '最近交易日': last_trading.strftime('%Y-%m-%d'),
            '当前时间': now.strftime('%H:%M:%S'),
        }
    
    @staticmethod
    def should_use_realtime_data() -> tuple:
        """
        判断是否应该使用实时数据
        
        Returns:
            tuple: (是否使用实时数据, 使用的日期)
        """
        today = datetime.date.today()
        
        # 如果不是交易日，使用最近的交易日数据
        if not TradingCalendar.is_trading_day(today):
            last_trading = TradingCalendar.get_last_trading_day(today)
            return False, last_trading
        
        # 如果是交易日，检查是否在交易时间内
        now = datetime.datetime.now()
        current_time = now.time()
        market_open_time = datetime.time(9, 30)
        
        # 如果还未开盘，使用前一个交易日数据
        if current_time < market_open_time:
            last_trading = TradingCalendar.get_last_trading_day(
                today - datetime.timedelta(days=1)
            )
            return False, last_trading
        
        # 开盘后使用实时数据
        return True, today
    
    @staticmethod
    def format_date_for_akshare(date: datetime.date) -> str:
        """
        将日期格式化为akshare需要的格式
        
        Args:
            date: 日期对象
            
        Returns:
            str: 格式化后的日期字符串 (YYYYMMDD)
        """
        return date.strftime('%Y%m%d')
    
    @staticmethod
    def get_next_trading_day(date: Optional[datetime.date] = None) -> datetime.date:
        """
        获取下一个交易日
        
        Args:
            date: 基准日期，默认为今天
            
        Returns:
            datetime.date: 下一个交易日
        """
        if date is None:
            date = datetime.date.today()
        
        # 往后查找下一个交易日（最多查找15天）
        for i in range(1, 16):
            check_date = date + datetime.timedelta(days=i)
            if TradingCalendar.is_trading_day(check_date):
                return check_date
        
        # 如果15天内都没找到，返回当前日期+1（异常情况）
        return date + datetime.timedelta(days=1)

